# -*- coding: utf-8 -*-
"""
批量测试LLM参数修改功能

测试各种参数修改请求，验证LLM是否能正确理解和返回JSON格式
"""
import sys
from pathlib import Path
from loguru import logger
import json
import time
import pandas as pd
from datetime import datetime

# 添加项目路径
project_root = Path(__file__).resolve().parent
sys.path.insert(0, str(project_root))

from services import ConversationManager

# 测试用例
TEST_CASES = [
    # 基本参数修改
    {
        "name": "修改K参数（使用参数名）",
        "message": "把k改成0.5",
        "expected_param": "parameters.K",
        "expected_value": 0.5,
        "should_succeed": True
    },
    {
        "name": "修改K参数（使用中文描述）",
        "message": "把表层出流系数改成0.6",
        "expected_param": "parameters.K",
        "expected_value": 0.6,
        "should_succeed": True
    },
    {
        "name": "修改U0参数（使用中文描述）",
        "message": "把下层蓄水容量改成50",
        "expected_param": "parameters.U0",
        "expected_value": 50,
        "should_succeed": True
    },
    {
        "name": "修改S0参数",
        "message": "把表层蓄水容量设置为30",
        "expected_param": "parameters.S0",
        "expected_value": 30,
        "should_succeed": True
    },
    {
        "name": "修改A参数",
        "message": "将蓄水空间分布参数调整为3.5",
        "expected_param": "parameters.A",
        "expected_value": 3.5,
        "should_succeed": True
    },
    {
        "name": "修改D0参数",
        "message": "把地下水蓄水容量改成100",
        "expected_param": "parameters.D0",
        "expected_value": 100,
        "should_succeed": True
    },
    # 初始状态修改
    {
        "name": "修改初始表层蓄水量",
        "message": "把初始表层蓄水量改成10",
        "expected_param": "initial_state.SA0",
        "expected_value": 10,
        "should_succeed": True
    },
    {
        "name": "修改初始下层蓄水量",
        "message": "将初始下层蓄水量设为20",
        "expected_param": "initial_state.UA0",
        "expected_value": 20,
        "should_succeed": True
    },
    # 多参数修改
    {
        "name": "同时修改多个参数",
        "message": "把K改成0.7，把A改成4.0",
        "expected_params": ["parameters.K", "parameters.A"],
        "expected_values": [0.7, 4.0],
        "should_succeed": True
    },
    # 边界值测试
    {
        "name": "修改到最小值",
        "message": "把K改成0.1",
        "expected_param": "parameters.K",
        "expected_value": 0.1,
        "should_succeed": True
    },
    {
        "name": "修改到最大值",
        "message": "把K改成0.9",
        "expected_param": "parameters.K",
        "expected_value": 0.9,
        "should_succeed": True
    },
    # 超出范围测试（应该拒绝修改）
    {
        "name": "K超出最小值（应拒绝）",
        "message": "把K改成0.05",
        "expected_param": "parameters.K",
        "expected_value": 0.05,
        "should_succeed": False,
        "reason": "值小于最小值0.1"
    },
    {
        "name": "K超出最大值（应拒绝）",
        "message": "把K改成1.5",
        "expected_param": "parameters.K",
        "expected_value": 1.5,
        "should_succeed": False,
        "reason": "值大于最大值0.9"
    },
    {
        "name": "A超出最大值（应拒绝）",
        "message": "把蓄水空间分布参数改成10",
        "expected_param": "parameters.A",
        "expected_value": 10,
        "should_succeed": False,
        "reason": "值大于最大值5"
    },
    {
        "name": "U0超出最大值（应拒绝）",
        "message": "把下层蓄水容量改成100",
        "expected_param": "parameters.U0",
        "expected_value": 100,
        "should_succeed": False,
        "reason": "值大于最大值90"
    },
    {
        "name": "S0超出最大值（应拒绝）",
        "message": "把表层蓄水容量改成60",
        "expected_param": "parameters.S0",
        "expected_value": 60,
        "should_succeed": False,
        "reason": "值大于最大值50"
    },
    {
        "name": "D0超出范围（应拒绝）",
        "message": "把地下水蓄水容量改成200",
        "expected_param": "parameters.D0",
        "expected_value": 200,
        "should_succeed": False,
        "reason": "值大于最大值160"
    },
    # 负数测试（应拒绝）
    {
        "name": "负数值（应拒绝）",
        "message": "把K改成-0.5",
        "expected_param": "parameters.K",
        "expected_value": -0.5,
        "should_succeed": False,
        "reason": "值小于最小值0.1"
    },
]

def run_test(manager, session_id, test_case):
    """运行单个测试用例"""
    logger.info(f"\n{'='*60}")
    logger.info(f"测试: {test_case['name']}")
    logger.info(f"用户消息: {test_case['message']}")
    if not test_case.get("should_succeed", True):
        logger.info(f"预期: 应拒绝修改 - {test_case.get('reason', '')}")
    logger.info(f"{'='*60}")
    
    try:
        # 记录开始时间
        start_time = time.time()
        
        # 发送消息
        response = manager.continue_conversation(
            session_id=session_id,
            user_message=test_case['message']
        )
        
        # 计算响应时间
        elapsed_time = time.time() - start_time
        
        logger.info(f"⏱️  响应时间: {elapsed_time:.2f}秒")
        
        if not response or not response.get("success"):
            error_msg = response.get('message') if response else '未知错误'
            logger.error(f"❌ 对话失败: {error_msg}")
            return {
                "passed": False,
                "elapsed_time": elapsed_time,
                "error": error_msg
            }
        
        # 检查结果
        understood = response.get("understood", False)
        parameter_changes = response.get("parameter_changes")
        message = response.get("message", "")
        explanation = response.get("explanation", "")
        validation_errors = response.get("validation_errors")
        
        logger.info(f"LLM回复: {message[:100]}...")
        logger.info(f"understood: {understood}")
        logger.info(f"parameter_changes: {parameter_changes}")
        logger.info(f"validation_errors: {validation_errors}")
        logger.info(f"explanation: {explanation[:100] if explanation else 'None'}...")
        
        should_succeed = test_case.get("should_succeed", True)
        
        # 对于应该拒绝的测试
        if not should_succeed:
            # 检查是否有验证错误（正确拒绝）
            if validation_errors:
                logger.info(f"✅ 正确拒绝: {validation_errors}")
                return {
                    "passed": True,
                    "elapsed_time": elapsed_time,
                    "error": None,
                    "llm_response": message,
                    "parameter_changes": parameter_changes,
                    "validation_errors": validation_errors
                }
            # 如果LLM理解了但没有修改（也算正确）
            elif understood and not parameter_changes:
                logger.info(f"✅ LLM拒绝修改")
                return {
                    "passed": True,
                    "elapsed_time": elapsed_time,
                    "error": None,
                    "llm_response": message,
                    "parameter_changes": None,
                    "validation_errors": None
                }
            # 如果LLM没理解（也算正确，因为没有错误修改）
            elif not understood:
                logger.info(f"✅ LLM未理解（未进行错误修改）")
                return {
                    "passed": True,
                    "elapsed_time": elapsed_time,
                    "error": None,
                    "llm_response": message,
                    "parameter_changes": None,
                    "validation_errors": None
                }
            # 如果成功修改了（错误！）
            else:
                logger.error(f"❌ 错误：应该拒绝但却修改了")
                return {
                    "passed": False,
                    "elapsed_time": elapsed_time,
                    "error": "应该拒绝但却修改了",
                    "llm_response": message,
                    "parameter_changes": parameter_changes,
                    "validation_errors": validation_errors
                }
        
        # 对于应该成功的测试
        if not understood:
            logger.warning(f"⚠️  LLM未理解请求")
            return {
                "passed": False,
                "elapsed_time": elapsed_time,
                "error": "LLM未理解请求",
                "llm_response": message,
                "parameter_changes": None,
                "validation_errors": None
            }
        
        if not parameter_changes:
            logger.warning(f"⚠️  没有参数修改")
            return {
                "passed": False,
                "elapsed_time": elapsed_time,
                "error": "没有参数修改",
                "llm_response": message,
                "parameter_changes": None,
                "validation_errors": None
            }
        
        # 检查期望的参数
        if "expected_param" in test_case:
            expected_param = test_case["expected_param"]
            expected_value = test_case["expected_value"]
            
            if expected_param in parameter_changes:
                actual_value = parameter_changes[expected_param]
                if actual_value == expected_value:
                    logger.info(f"✅ 测试通过: {expected_param} = {actual_value}")
                    return {
                        "passed": True,
                        "elapsed_time": elapsed_time,
                        "error": None,
                        "llm_response": message,
                        "parameter_changes": parameter_changes,
                        "validation_errors": validation_errors
                    }
                else:
                    error_msg = f"值不匹配: 期望 {expected_value}, 实际 {actual_value}"
                    logger.warning(f"⚠️  {error_msg}")
                    return {
                        "passed": False,
                        "elapsed_time": elapsed_time,
                        "error": error_msg,
                        "llm_response": message,
                        "parameter_changes": parameter_changes,
                        "validation_errors": validation_errors
                    }
            else:
                error_msg = f"未找到参数: {expected_param}"
                logger.warning(f"⚠️  {error_msg}")
                logger.info(f"实际参数: {list(parameter_changes.keys())}")
                return {
                    "passed": False,
                    "elapsed_time": elapsed_time,
                    "error": error_msg,
                    "llm_response": message,
                    "parameter_changes": parameter_changes,
                    "validation_errors": validation_errors
                }
        
        elif "expected_params" in test_case:
            expected_params = test_case["expected_params"]
            expected_values = test_case["expected_values"]
            
            all_match = True
            errors = []
            for param, value in zip(expected_params, expected_values):
                if param in parameter_changes:
                    actual_value = parameter_changes[param]
                    if actual_value == value:
                        logger.info(f"✅ {param} = {actual_value}")
                    else:
                        error_msg = f"{param}: 期望 {value}, 实际 {actual_value}"
                        logger.warning(f"⚠️  {error_msg}")
                        errors.append(error_msg)
                        all_match = False
                else:
                    error_msg = f"未找到参数: {param}"
                    logger.warning(f"⚠️  {error_msg}")
                    errors.append(error_msg)
                    all_match = False
            
            return {
                "passed": all_match,
                "elapsed_time": elapsed_time,
                "error": "; ".join(errors) if errors else None,
                "llm_response": message,
                "parameter_changes": parameter_changes,
                "validation_errors": validation_errors
            }
        
        return {
            "passed": True,
            "elapsed_time": elapsed_time,
            "error": None,
            "llm_response": message,
            "parameter_changes": parameter_changes,
            "validation_errors": validation_errors
        }
        
    except Exception as e:
        logger.error(f"❌ 测试异常: {str(e)}", exc_info=True)
        return {
            "passed": False,
            "elapsed_time": 0,
            "error": str(e),
            "llm_response": "",
            "parameter_changes": None,
            "validation_errors": None
        }


def main():
    """主函数"""
    logger.info("=" * 60)
    logger.info("批量测试LLM参数修改功能")
    logger.info("=" * 60)
    
    # 初始化对话管理器
    manager = ConversationManager()
    
    # 使用实际的session_id
    session_id = "de9a990b-59e2-466a-84af-1e0bc1f3c585"
    
    logger.info(f"\n使用会话ID: {session_id}")
    
    # 运行所有测试
    results = []  # 用于CSV输出
    test_stats = []  # 用于统计
    total_time = 0
    
    for i, test_case in enumerate(TEST_CASES, 1):
        logger.info(f"\n\n进度: {i}/{len(TEST_CASES)}")
        result = run_test(manager, session_id, test_case)
        
        # 构建预期输出描述
        should_succeed = test_case.get("should_succeed", True)
        if should_succeed:
            if "expected_param" in test_case:
                expected_output = f"修改 {test_case['expected_param']} = {test_case['expected_value']}"
            elif "expected_params" in test_case:
                changes = [f"{p}={v}" for p, v in zip(test_case['expected_params'], test_case['expected_values'])]
                expected_output = f"修改 {', '.join(changes)}"
            else:
                expected_output = "成功修改参数"
        else:
            expected_output = f"拒绝修改（{test_case.get('reason', '超出范围')}）"
        
        # 构建模型输出描述
        if result.get("validation_errors"):
            model_output = f"拒绝: {result['validation_errors']}"
        elif result.get("parameter_changes"):
            model_output = f"修改: {json.dumps(result['parameter_changes'], ensure_ascii=False)}"
        elif result.get("llm_response"):
            model_output = result["llm_response"][:100]
        else:
            model_output = "未理解或拒绝"
        
        # 保存CSV格式的结果
        results.append({
            "用户输入": test_case["message"],
            "预期输出": expected_output,
            "模型输出": model_output,
            "模型返回是否正确": "是" if result["passed"] else "否",
            "运行时间(秒)": round(result["elapsed_time"], 2)
        })
        
        # 保存统计用的数据
        test_stats.append({
            "name": test_case["name"],
            "passed": result["passed"],
            "elapsed_time": result["elapsed_time"],
            "error": result.get("error")
        })
        
        total_time += result["elapsed_time"]
    
    # 输出总结
    logger.info("\n\n" + "=" * 60)
    logger.info("测试结果汇总")
    logger.info("=" * 60)
    
    passed_count = sum(1 for r in test_stats if r["passed"])
    total_count = len(test_stats)
    
    for stat in test_stats:
        status = "✅ 通过" if stat["passed"] else "❌ 失败"
        time_str = f"({stat['elapsed_time']:.2f}s)"
        error_str = f" - {stat['error']}" if stat.get('error') else ""
        logger.info(f"{status} {time_str} - {stat['name']}{error_str}")
    
    logger.info(f"\n总计: {passed_count}/{total_count} 测试通过")
    logger.info(f"通过率: {passed_count/total_count*100:.1f}%")
    logger.info(f"总耗时: {total_time:.2f}秒")
    logger.info(f"平均响应时间: {total_time/total_count:.2f}秒")
    
    if passed_count == total_count:
        logger.info("\n🎉 所有测试通过！")
    else:
        logger.warning(f"\n⚠️  {total_count - passed_count} 个测试失败")
    
    # 保存结果到CSV
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    csv_filename = f"test_results_{timestamp}.csv"
    csv_path = Path("storage") / csv_filename
    
    # 创建DataFrame
    df = pd.DataFrame(results)
    
    # 保存到CSV（使用utf-8-sig编码以支持Excel正确显示中文）
    df.to_csv(csv_path, index=False, encoding='utf-8-sig')
    
    logger.info(f"\n📊 测试结果已保存到: {csv_path}")
    logger.info(f"   列: 用户输入 | 预期输出 | 模型输出 | 模型返回是否正确 | 运行时间(秒)")
    logger.info("=" * 60)
    
    return passed_count == total_count


if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)
